package org.covalence.core.instrument;

import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;

import org.apache.log4j.Logger;
import org.covalence.core.Registry;
import org.covalence.core.visitor.BoundClassVisitor;
import org.covalence.core.visitor.ClassConfigVisitor;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;

public class EnhancingAgent implements ClassFileTransformer {

	private static Logger log = Logger.getLogger(EnhancingAgent.class);

	public byte[] transform(ClassLoader loader, String className,
			Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
			byte[] bytes) throws IllegalClassFormatException {

		if (Registry.shouldEnhance(className)) {	
			ClassWriter cw = new ClassWriter(true);
			BoundClassVisitor bma = new BoundClassVisitor(cw);
			
			ClassReader cr = new ClassReader(bytes);
			ClassConfigVisitor cfr = new ClassConfigVisitor(className);
			
			log.debug("Examining class: " + className);
			cr.accept(cfr, true);
			Registry.register(cfr.getClassConfig());	
			
			log.debug("Transforming class: " + className);
			cr.accept(bma, false);
			
			byte[] results = cw.toByteArray();
			
			log.debug("Attempting to write: " + className + " to file.");
			_save(className, results);
			
			return results;
		} else {
			return null;
		}
	}

    private void _save(String className, byte[] b) {
    	OutputStream out = null;
    	try {
    		String dirs = className.replace('.', File.separatorChar);
    		new File("debug" + File.separatorChar + dirs).getParentFile().mkdirs();
    		File file = new File(new File("debug"), dirs + "$ASM$.class");
    		log.debug("Writing file: " + file.toString());
    		out = new BufferedOutputStream(new FileOutputStream(file));
            out.write(b);
        } catch (IOException e) {
			e.printStackTrace();
		} finally {
            try {
				out.close();
				log.debug("Wrote and closed the file.");
			} catch (IOException e) {
				e.printStackTrace();
			}
        }
    }

	public static void premain(String arglist, Instrumentation inst) throws ClassNotFoundException, UnmodifiableClassException {
        inst.addTransformer(new EnhancingAgent());
    }
}
